#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import threading

import tyrant
from tyrant import Query, QueryCondition, QueryOrder, IndexType


#--- Connection manager ----------------------------------------------
connections = threading.local()

def get_connection(host, port):
    key = '%s:%s' % (host, port)
    cur_connection = hasattr(connections, key)

    if not cur_connection:
        setattr(connections, key, tyrant.Table.open(host, port))
    return getattr(connections, key)

def close_open_connections():
    for key in dir(connections):
        if '__' not in key:
            getattr(connections, key).close()
            delattr(connections, key)

def get_open_connections():
    for key in dir(connections):
        if '__' not in key:
          yield getattr(connections, key)

#--- pytyrant client ----------------------------------------------
class TyrantClient:
    def __init__(self, servers):
        to_int = lambda item: (item[0], int(item[1]))
        self.servers = [ to_int(s.split(':')) for s in servers ]

    #--- Incr and decr ----------------------------------------------
    #def incr(self, key, delta=1):
        #db = self.get_db(key)
        #return db.call_func('incr', key, "%s" % delta, record_locking=True)

    #--- Set, get and delete ----------------------------------------------
    def set(self, key, val, replace=False, **kw):
        db = self.get_db(key)

        try:
            if replace:
                ret = db.put(key, val)
            else:
                ret = db.putcat(key, val)
            return ret
        except:
            return False

    #TODO: add mget
    def get(self, key, **kw):
        db = self.get_db(key)

        try:
            val = db.get(key)
        except:
            return None
        return val

    def delete(self, key):
        db = self.get_db(key)

        try:
            db.out(key)
            return True
        except:
            return False

    #--- Query ---------------------------------------------
    def query(self, query, attributes=None):
        db = self.get_db(query)

        try:
            query.rdb = db
            if attributes:
                if isinstance(attributes, bool):
                    return query.search_get()
                else:
                    return query.search_get(attributes)
            else:
                return query.search()
        except:
            return []

    #--- Set index -----------------------------------------
    def set_index(self, name, type):
        db = self.get_db(name)

        try:
            return db.setindex(name, type)
        except:
            return False


    #--- List ----------------------------------------------
    #def _encode_list(self, values):
        #v_encoded = []
        #for v in values:
            #v_encoded.append('%s~' % v)
        #return ''.join(v_encoded)

    #def list_add(self, key, values):
        #db = self.get_db(key)
        #return db.call_func('list_add', key, self._encode_list(values), record_locking=True)

    #def list_remove(self, key, values):
        #db = self.get_db(key)
        #return db.call_func('list_remove', key, self._encode_list(values), record_locking=True)

    def get_db(self, key):
        servers = list(self.servers)

        #Load balance by key
        index = hash(key) % len(servers)
        first_host, first_port = servers.pop(index)

        try:
            return get_connection(first_host, first_port)
        except:
            pass

        #This did not work out, try the other servers
        for host, port in servers:
            try:
                return get_connection(host, port)
            except Exception, e:
                if e[0] == 61:
                    continue
                raise
        raise


class TyrantNode(TyrantClient):
    """Extends the tyrant client with a proper __str__ method"""

    def __init__(self, name, nodes, *k, **kw):
        self.name = name
        TyrantClient.__init__(self, nodes, *k, **kw)

    def __str__(self):
        return self.name
